from core.function import Function
from core.cuda import cuda


class Square(Function):
    def forward(self, x):
        """
        执行 Square 操作的前向传播。

        Parameters:
            x (Variable): 输入变量。

        Returns:
            Variable: Square 操作的结果。
        """
        x = cuda.to_array(x)
        y = x ** 2
        return y

    def backward(self, gy):
        """
        执行 Square 操作的反向传播。

        Parameters:
            gy (Variable): 相对于前向传播结果的梯度。

        Returns:
            Variable: 相对于输入变量的梯度。
        """

        # 多个输入
        x, = self.inputs
        gx = 2 * x * gy
        return gx

    def square(self, x):
        """
        执行 Square 操作。

        Parameters:
            x (Variable): 输入变量。

        Returns:
            Variable: Square 操作的结果。
        """
        return self(x)

